from typing import Annotated

from IPython.core.display import Image
from IPython.core.display_functions import display
from langchain_anthropic import ChatAnthropic
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode, tools_condition
from typing_extensions import TypedDict

from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages

from utils.llms import get_deepseek


class State(TypedDict):
    messages: Annotated[list, add_messages]


graph_builder = StateGraph(State)


@tool
def search_tool(time: str = None):
    """
        搜索航班
        Args:
        time: 要查询的时间
    """
    print("Searching...", time)
    return [
        {
            "title": "Flight 1",
            "description": "This flight is from New York to Los Angeles.",
        },
        {
            "title": "Flight 2",
            "description": "This flight is from Los Angeles to New York.",
        },
    ]


tools = [search_tool]
# tool.invoke("What's a 'node' in LangGraph?")

llm = get_deepseek()
llm_with_tools = llm.bind_tools(tools)


def chatbot(state: State):
    return {"messages": [llm_with_tools.invoke(state["messages"])]}


graph_builder.add_node("chatbot", chatbot)

tool_node = ToolNode(tools=tools)
graph_builder.add_node("tools", tool_node)

graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")
graph_builder.set_entry_point("chatbot")
graph = graph_builder.compile()

graph.get_graph().draw_mermaid_png(output_file_path="./graph.png")


def stream_graph_updates(user_input: str):
    for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)


while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break
